import numpy as np
import argparse
import cv2

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='show npy image')
    parser.add_argument('npy_path', default=None, metavar='NPY_PATH', type=str,
                        help='Path to the npy file')
    parser.add_argument('-o', '--saveas', default=None, metavar='IMG_PATH', type=str,
                        help='Path to the image file, or - output as [npy_path].jpg')
    args = parser.parse_args()

    npy = np.load(args.npy_path)
    if npy.dtype == "float64":
        npy = np.uint8(npy * 255)
        npy = cv2.applyColorMap(npy, cv2.COLORMAP_JET)

    if args.saveas:
        if args.saveas == '-':
            args.saveas = args.npy_path + '.jpg'
        cv2.imwrite(args.saveas, npy)
    
    cv2.namedWindow(args.saveas, cv2.WINDOW_NORMAL)
    cv2.imshow(args.saveas, npy)
    cv2.waitKey(0)
    cv2.destroyAllWindows()